import numpy as np
from sklearn.model_selection import train_test_split
from bishe_situations.utils import GPLEARN_PARAMS, OPS_GPLEARN, generate_data, BaselineTrainer
from gplearn.genetic import SymbolicRegressor

class GpLearnTrainer(BaselineTrainer):
    def __init__(self, params: dict = GPLEARN_PARAMS, population_size: int = 50, generations: int = 20, **kwargs):
        super().__init__(params, population_size, generations, **kwargs)
        self.model = SymbolicRegressor(
            function_set=OPS_GPLEARN,
            verbose=1,
            random_state=42,
            population_size=population_size,
            generations=generations,
            metric='mse',
            **params,
            **kwargs
        )

    def override_fit(self, X, y):
        X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)
        self.model.fit(X_train, y_train.ravel())
        self.mse = self.model.score(X_val, y_val.ravel())

    def best_function(self):
        return str(self.model._program)
    

def optimize_hyperparameters(X, y, n_trials=10):

    model = GpLearnTrainer()
    model.fit(X, y)
    
    return model.best_function(), model.mse, model.train_time

def main():
    # 定义目标函数
    def target_function(X):
        return np.sin(X) * 2 + np.cos(X) - 80 + np.exp(X)
    
    # 生成数据
    X, y = generate_data(target_function, n_samples=512, noise=0.8)
    
    # 优化超参数
    print("开始优化超参数...")
    best_equation, best_mse, best_train_time = optimize_hyperparameters(X, y)
    
    print("\n找到的最佳公式:")
    print(f"y = {best_equation}")
    print(f"\nMSE: {best_mse:.4f}")
    print(f"\n训练时间: {best_train_time:.4f}秒")

if __name__ == "__main__":
    main() 